#!/bin/tcsh -ef
####
#
#############################################################################
#                                                                           #
# Title: Align w/ MotionCor2                                                #
#                                                                           #
# (C) 2dx.org, GNU Plublic License.                                         #
#                                                                           #
# Created..........: 07/21/2016                                             #
# Last Modification: 07/21/2016                                             #
# Author...........: 2dx.org                                                #
#                                                                           #
#############################################################################
#
# SORTORDER: 20 
#
# MANUAL: This script runs a drift correction program on a movie file, to produce a drift-corrected movie and to produce a drift-corrected and averaged image file.
#
# PUBLICATION:<A HREF="http://msg.ucsf.edu/em/software/motioncor2.html">MotionCor2 Website</A>
#
# DISPLAY: imagenumber
# DISPLAY: imagename_original
# DISPLAY: comment
# DISPLAY: sample_pixel
# DISPLAY: movie_imagenumber_total
# DISPLAY: KV
# DISPLAY: SERIALEM_FACTOR
# DISPLAY: total_dose
# DISPLAY: frame_dose
# DISPLAY: generate_dose_dw
# DISPLAY: apply_dw
# DISPLAY: num_frames_sum
# DISPLAY: movie_stackname
# DISPLAY: MotionCor2_throw
# DISPLAY: MotionCor2_param1
# DISPLAY: MotionCor2_param2
# DISPLAY: MotionCor2_param3
# DISPLAY: import_bin
# DISPLAY: import_bin_target
# DISPLAY: import_drift
# DISPLAY: import_drift_longest
# DISPLAY: import_drift_deceleration
# DISPLAY: import_drift_jitter
# DISPLAY: tempkeep
# DISPLAY: do_unpack
# DISPLAY: import_gainref
#
#$end_local_vars
#
set bin_2dx = ""
set proc_2dx = ""
set app_2dx_mrc_converter = ""
#
set tempkeep = ""
set imagename = ""
set nonmaskimagename = ""
set imagenumber = ""
set imagesidelength = ""
set sample_pixel = ""
set raw_gaincorrectedstack = ""
set movie_stackname = ""
set movie_imagenumber_total = ""
set import_rawstack = ""
set import_gainref = ""
set KV = ""
set SERIALEM_FACTOR = ""
set total_dose = ""
set frame_dose = ""
set import_bin = ""
set import_bin_target = ""
set import_original_time = ""
set import_gainref = ""
set MotionCor2_param1 = ""
set MotionCor2_param2 = ""
set MotionCor2_param3 = ""
set import_drift = ""
set import_drift_longest = ""
set import_drift_deceleration = ""
set import_drift_jitter = ""
set generate_dose_dw = ""
set apply_dw = ""
set num_frames_sum = ""
set MotionCor2_throw = ""
set do_unpack = ""
#
#$end_vars
#
#NOT USED AT THE MOMENT:
# set MotionCor2_inframemotion = ""
# DISPLAY: import_MotionCor2_DW_option
# set import_MotionCor2_DW_option = ""
# DISPLAY: MotionCor2_delete_non_driftcor_stack
# set MotionCor2_delete_non_driftcor_stack = ""
#
# get motioncor2 from
# wget http://msg.ucsf.edu/MotionCor2/MotionCor2.zip
#
set scriptname = import_driftcor_motioncor2
\rm -f LOGS/${scriptname}.results
#
source ${proc_2dx}/initialize
#
source ${proc_2dx}/2dx_makedirs
#
if ( ${website_blur_doit} == "y" ) then
  set blurkernel = ${website_blurkernel}
else
  set blurkernel = 0
endif
#
if(${import_original_time} == "-" || ${import_original_time} == "") then
  @ status_date = `date +%s` * 1000
  set date_text = "Processed at "`echo ${status_date} | awk '{ s = $1 / 1000 } END { print s }' | ${app_gawk} '{print strftime("%c", $0)}' `
else
  set status_date = ${import_original_time}
  set date_text = "Recorded at "`echo ${status_date} | awk '{ s = $1 / 1000 } END { print s }' | ${app_gawk} '{print strftime("%c", $0)}' `
endif
#
echo "<<@evaluate>>"
#
echo "<<@progress: 1>>"
#
set framedose_too_small = `echo ${frame_dose} | awk '{ if ( $1 < 0.001 ) { s = 1 } else { s = 0 } } END { print s }'`
if ( ${framedose_too_small} == 1 ) then
  set frame_dose = 0.001
  echo ":: "
  ${proc_2dx}/linblock "WARNING: frame_dose corrected to ${frame_dose}"
  echo ":: "
endif
#
if ( ${movie_stackname} == "ScriptWillPutNameHere" ) then
  set movie_stackname = movie_aligned
  echo "set movie_stackname = ${movie_stackname}" >> LOGS/${scriptname}.results
endif
#
echo "# IMAGE-IMPORTANT: ${raw_gaincorrectedstack}.mrc <Gain-corrected average image (2D)>" >> LOGS/${scriptname}.results
echo "# IMAGE-IMPORTANT: ${raw_gaincorrectedstack}_fft.mrc <Gain-corrected average image FFT (2D)>" >> LOGS/${scriptname}.results
#
if ( ${do_unpack} == "y" ) then
  if ( ! -e FASTDISK/${raw_gaincorrectedstack}.mrcs ) then
    ${proc_2dx}/protest "ERROR: Input movie FASTDISK/${raw_gaincorrectedstack}.mrcs not found."
  endif
endif
#
echo ": "
${proc_2dx}/lin "GPU_how_many = ${GPU_how_many}"
${proc_2dx}/lin "GPU_to_use = ${GPU_to_use}"
${proc_2dx}/lin "GPU_cycle = ${GPU_cycle}"
echo ": "
#

#
##########################################################################
${proc_2dx}/linblock "Calling MotionCor2..."
##########################################################################
#
if ( ! -e ${app_motioncor2} ) then
  ${proc_2dx}/protest "ERROR: First define location of MOTIONCOR2 in Settings."
endif
#
\rm -f ${movie_stackname}.mrcs
\rm -f ${movie_stackname}_motioncor2Sum.mrc
\rm -f ${movie_stackname}_motioncor2Sum_DW.mrc
\rm -f ${movie_stackname}_motioncor2Sum_Stk.mrc
#
echo "<<@progress: 20>>"
echo "<<@evaluate>>"
#
if ( ${import_bin} == "1" ) then
  ${proc_2dx}/linblock "Fourier-cropping the image, i.e., binning by 2, after drift correction."
  set motioncorparam1 = "-OutStack 1 -ftbin 2 "
else
  set motioncorparam1 = "-OutStack 1 "
endif
#
echo "# IMAGE: LOGS/motioncor2.out <MotionCor2 Output (TXT)>" >> LOGS/${scriptname}.results
#
set olddir = $PWD
cd FASTDISK
#
if ( ${GPU_how_many} > "1" ) then
  if ( ${GPU_cycle} == "y" ) then
    #source ${proc_2dx}/2dx_find_GPU.com
    set GPU_to_use_local = "${next_GPU}"
  else
    set GPU_to_use_local = "${GPU_to_use}" 
  endif
else
  set GPU_to_use_local = "${GPU_to_use}"
endif
#
if ( ${do_unpack} == "y" ) then
  set MC2_input = "-InMrc ${raw_gaincorrectedstack}.mrcs"
else
  set import_gainref_noending = `echo "${import_gainref}" | rev | cut -d\. -f2- | rev`
  if ( ! -e "${olddir}/${import_gainref_noending}.mrc" ) then
    if ( -e "${olddir}/${import_gainref_noending}.MRC" ) then
      mv "${olddir}/${import_gainref_noending}.MRC" "${olddir}/${import_gainref_noending}.mrc"
    else 
      if ( -e "${olddir}/${import_gainref_noending}.dm4" ) then
        dm2mrc "${olddir}/${import_gainref_noending}.dm4" "${olddir}/${import_gainref_noending}.mrc"
      else 
        ${proc_2dx}/protest "ERROR: Gainreference ${olddir}/${import_gainref_noending}.dm4 not found."
      endif
    endif
  endif
  set MC2_input = "-InTiff ${olddir}/${import_rawstack} -Gain ${olddir}/${import_gainref_noending}.mrc"
endif
#
echo ": "
echo ": "${app_motioncor2} 
echo ": " ${MC2_input}
echo ": " -OutMrc ${movie_stackname}_motioncor2Sum.mrc
echo ": " -Iter 15
echo ": " -Tol 0.4
echo ": " -kV ${KV}
echo ": " -PixSize ${sample_pixel}
echo ": " -FmDose ${frame_dose}
echo ": " -Gpu ${GPU_to_use_local}  
echo ": " -Throw ${MotionCor2_throw}
echo ": " ${motioncorparam1}
echo ": " ${MotionCor2_param1}
echo ": " ${MotionCor2_param2}
echo ": " ${MotionCor2_param3}
#echo ": " -InFmMotion 1 # ${MotionCor2_inframemotion}
echo ": "
#

if ( ${MotionCor2_alone} == "y" ) then
  # set small random delay < 1s to prevent multiple motioncor2 processes from being spawned at exactly the same time
  set i = 0
  while ( ${i} < 1 )
    set sleeptime = `bash -c 'echo $RANDOM % 50 | bc'`
    set sleeptime = `echo "scale=4; ${sleeptime} / 100" | bc | tr -d "\n"`
    sleep ${sleeptime}s
    @ i++ 
  end
  ps -al > tmp.1
  set isrunning = `cat tmp.1 | grep -i MotionCor2 | wc -l`
  set icount = 1
  while ( ${isrunning} != '0' ) 
    set sleeptime = `bash -c 'echo $RANDOM % 5 + 1 | bc'`
    echo "::MotionCor2 appears to be running. Waiting (${icount}) x time, for ${sleeptime} seconds..."
    sleep ${sleeptime}
    ps -al > tmp.1
    set isrunning = `cat tmp.1 | grep -i MotionCor2 | wc -l`
    @ icount += 1
    if ( ${icount} >= 300 ) then
      ${proc_2dx}/protest "ERROR: MotionCor2 seems to be running endlessly. Something must be wrong."
    endif
  end
  \rm -f tmp.1
endif

time ${app_motioncor2} \
${MC2_input} \
-OutMrc ${movie_stackname}_motioncor2Sum.mrc \
-Iter 15 \
-Tol 0.4 \
-kV ${KV} \
-PixSize ${sample_pixel} \
-FmDose ${frame_dose} \
-Gpu ${GPU_to_use_local}   \
-Throw ${MotionCor2_throw} \
${motioncorparam1} \
${MotionCor2_param1} \
${MotionCor2_param2} \
${MotionCor2_param3} \
 > ${olddir}/LOGS/motioncor2.out
#
#-InFmMotion ${MotionCor2_inframemotion} \
#
\rm -f tmp.1
\rm -f tmp.2

# Call Python script to generate a DW aligned average with frames only up to desired dose:
if ( ${generate_dose_dw} == "y" ) then
	echo ": Now generating DW aligned-average up to desired total dose:"
	echo ": ${app_python} ${proc_2dx}/apply_dose_filter.py ${movie_stackname}_motioncor2Sum_Stk.mrc ${movie_stackname}_motioncor2Sum_DW-rec.mrc ${sample_pixel} ${frame_dose} ${MotionCor2_throw} ${num_frames_sum} ${KV} ${apply_dw}"
	time ${app_python} ${proc_2dx}/apply_dose_filter.py ${movie_stackname}_motioncor2Sum_Stk.mrc ${movie_stackname}_motioncor2Sum_DW-rec.mrc ${sample_pixel} ${frame_dose} ${MotionCor2_throw} ${num_frames_sum} ${KV} ${apply_dw}
endif
#
cd ${olddir}
#
\mv -f FASTDISK/${movie_stackname}_motioncor2Sum_Stk.mrc ${movie_stackname}.mrcs
echo "# IMAGE-IMPORTANT: ${movie_stackname}.mrcs <Drift-corrected movie (stack)>" >> LOGS/${scriptname}.results
#
if ( ${import_bin} == "1" ) then
  set sample_pixel = `echo ${sample_pixel} | awk '{ s = 2 * $1 } END { print s }'`
  echo "set sample_pixel = ${sample_pixel}"  >> LOGS/${scriptname}.results
endif
#
cat LOGS/motioncor2.out
#
cat LOGS/motioncor2.out | grep shift: | grep Frame | cut -d\: -f2 > motioncor2_shifts.txt
#
if ( ! -e FASTDISK/${movie_stackname}_motioncor2Sum.mrc ) then
  ${proc_2dx}/protest "ERROR: ${movie_stackname}_motioncor2Sum.mrc was not created.  Aborting."
endif
#
#################################################################################
${proc_2dx}/linblock "Running: labelh.exe to normalize MRC image"
#################################################################################
#
echo "<<@progress: 45>>"
#
\rm -f ${movie_stackname}_Sum.mrc
#
time ${bin_2dx}/labelh.exe << eot
FASTDISK/${movie_stackname}_motioncor2Sum.mrc
42
${movie_stackname}_Sum.mrc
eot
#
\rm -f FASTDISK/${movie_stackname}_motioncor2Sum.mrc
#
echo "# IMAGE-IMPORTANT: ${movie_stackname}_Sum.mrc <Drift-corrected average (MRC, no DW)>" >> LOGS/${scriptname}.results
#
##########################################################################
${proc_2dx}/linblock "Running ${proc_2dx}/motioncor2_plotter.py motioncor2_shifts.txt translations.png drift_plotter.out ${sample_pixel}"
##########################################################################
#
echo "<<@progress: 50>>"
#
\rm -f translations.png
\rm -f drift_plotter.out
echo ":: app_python = ${app_python}"
echo "PWD = "$PWD
${app_python}  ${proc_2dx}/motioncor2_plotter.py motioncor2_shifts.txt translations.png drift_plotter.out ${sample_pixel}
#
source drift_plotter.out
echo "set import_drift = ${import_drift}" >> LOGS/${scriptname}.results
echo "set import_drift_longest = ${import_drift_longest}" >> LOGS/${scriptname}.results
echo "set import_drift_deceleration = ${import_drift_deceleration}" >> LOGS/${scriptname}.results
echo "set import_drift_jitter = ${import_drift_jitter}" >> LOGS/${scriptname}.results
#
echo "# IMAGE-IMPORTANT: translations.png <drift plot (PNG)>" >> LOGS/${scriptname}.results
#
if ( ${do_unpack} == "y" ) then
  \mv -f FASTDISK/${raw_gaincorrectedstack}.mrcs .
  echo "# IMAGE-IMPORTANT: ${raw_gaincorrectedstack}.mrcs <Raw movie (stack)>" >> LOGS/${scriptname}.results
endif
#
#################################################################################
${proc_2dx}/linblock "Running: labelh.exe to normalize MRC image"
#################################################################################
#
echo "<<@progress: 70>>"
#
\rm -f ${movie_stackname}.mrc
#
${bin_2dx}/labelh.exe << eot
FASTDISK/${movie_stackname}_motioncor2Sum_DW.mrc
42
${movie_stackname}.mrc
eot
#
\rm -f FASTDISK/${movie_stackname}_motioncor2Sum_DW.mrc
#
echo "# IMAGE-IMPORTANT: ${movie_stackname}.mrc <Drift-corrected average image (2D, DoseWeighted)>" >> LOGS/${scriptname}.results
#
# if ( -e FASTDISK/${movie_stackname}_motioncor2Sum_DW-rec.mrc ) then
# 	echo ":: CHECKPOINT "
# 	#################################################################################
# 	${proc_2dx}/linblock "Running: labelh.exe to normalize MRC image"
# 	#################################################################################
# 	#
# 	echo "<<@progress: 70>>"
# 	#
# 	\rm -f ${movie_stackname}_DW-rec.mrc
# 	#
# 	${bin_2dx}/labelh.exe << eot
# 	FASTDISK/${movie_stackname}_motioncor2Sum_DW-rec.mrc
# 	42
# 	${movie_stackname}_DW-rec.mrc
# 	eot
# 	#
	if ( -e FASTDISK/${movie_stackname}_motioncor2Sum_DW-rec.mrc ) then
		\mv FASTDISK/${movie_stackname}_motioncor2Sum_DW-rec.mrc ${movie_stackname}_DW-rec.mrc
	# 	#
		echo "# IMAGE-IMPORTANT: ${movie_stackname}_DW-rec.mrc <Drift-corrected average image (2D, DoseWeighted-Rec)>" >> LOGS/${scriptname}.results
	# 	#
	endif
# endif
#################################################################################
${proc_2dx}/linblock "Running: 2dx_mrc_converter to create a PNG file for inspection"
#################################################################################
#
echo "<<@progress: 80>>"
#
\rm -f ${movie_stackname}.mrc.png
\rm -f tmp.png
\rm -f tmp2.png
\rm -f STATUS/2-image.jpg
${app_2dx_mrc_converter} --size 400 ${movie_stackname}.mrc tmp.png 
${app_python} ${proc_2dx}/PNGannotator.py tmp.png tmp2.png 10 350 0 "Drift-corrected average (DW)"
${app_python} ${proc_2dx}/PNGannotator.py tmp2.png ${movie_stackname}.mrc.png 10 365 0 "${date_text}"
if ( ${blurkernel} != "0" ) then
  ${app_python} ${proc_2dx}/PNGannotator.py tmp.png tmp2.png 135 180 ${blurkernel} "I N T E N T I O N A L L Y"
  ${app_python} ${proc_2dx}/PNGannotator.py tmp2.png tmp.png 135 195 0             "   B  L  U  R  R  E  D"
endif
${app_python} ${proc_2dx}/PNGannotator.py tmp.png tmp2.png 10 350 0 "Drift-corrected average (DW)"
${app_python} ${proc_2dx}/PNGannotator.py tmp2.png STATUS/2-image.jpg 10 365 0 "${date_text}"
\rm -f tmp.png
\rm -f tmp2.png
#
##########################################################################
echo "::Running: clip fft ${movie_stackname}.mrc ${movie_stackname}_fft.mrc"
##########################################################################
#
echo "<<@progress: 90>>"
#
\rm -f ${movie_stackname}_fft.mrc
time nice +19 ${dir_imod}/bin/clip fft ${movie_stackname}.mrc ${movie_stackname}_fft.mrc
echo "# IMAGE-IMPORTANT: ${movie_stackname}_fft.mrc <Drift-corrected average image FFT (2D, DW)>" >> LOGS/${scriptname}.results
#
if ( -e FASTDISK/${movie_stackname}_motioncor2Sum_DW-rec.mrc ) then
	#
	\rm -f ${movie_stackname}_DW-rec_fft.mrc
	time nice +19 ${dir_imod}/bin/clip fft ${movie_stackname}_DW-rec.mrc ${movie_stackname}_DW-rec_fft.mrc
	echo "# IMAGE-IMPORTANT: ${movie_stackname}_DW-rec_fft.mrc <Drift-corrected average image FFT (2D, DW-rec)>" >> LOGS/${scriptname}.results
	#
endif
\rm -f ${movie_stackname}_Sum_fft.mrc
time nice +19 ${dir_imod}/bin/clip fft ${movie_stackname}_Sum.mrc ${movie_stackname}_Sum_fft.mrc
echo "# IMAGE-IMPORTANT: ${movie_stackname}_Sum_fft.mrc <Drift-corrected average image FFT (2D, non-DW)>" >> LOGS/${scriptname}.results
#
\rm -f ${movie_stackname}_fft.mrc.png
\rm -f tmp.png
\rm -f tmp2.png
${app_2dx_mrc_converter} --size 400 ${movie_stackname}_fft.mrc tmp.png 
${app_python} ${proc_2dx}/PNGannotator.py tmp.png tmp2.png 10 350 0 "FFT of drift-corrected average (DW)"
${app_python} ${proc_2dx}/PNGannotator.py tmp2.png ${movie_stackname}_fft.mrc.png 10 365 0 "${date_text}"
#
\rm -f ${movie_stackname}_Sum_fft.mrc.png
\rm -f tmp.png
\rm -f tmp2.png
\rm -f STATUS/2-FFT.jpg
${app_2dx_mrc_converter} --size 400 ${movie_stackname}_Sum_fft.mrc tmp.png 
${app_python} ${proc_2dx}/PNGannotator.py tmp.png tmp2.png 10 350 0 "FFT of drift-corrected average (non-DW)"
${app_python} ${proc_2dx}/PNGannotator.py tmp2.png ${movie_stackname}_Sum_fft.mrc.png 10 365 0 "${date_text}"
${app_python} ${proc_2dx}/PNGannotator.py tmp2.png STATUS/2-FFT.jpg 10 365 0 "${date_text}"
\rm -f tmp.png
\rm -f tmp2.png
#
set movie_enable = "y"
echo "set movie_enable = ${movie_enable}" >> LOGS/${scriptname}.results
#
##########################################################################
${proc_2dx}/linblock "${scriptname} - normal end."
##########################################################################
#
echo "<<@progress: 100>>"
echo "<<@evaluate>>"
exit
#
# These are listed here to make sure they appear in the 2dx_image GUI:
${app_python} ${proc_2dx}/apply_dose_filter.py
#
